package com.caishi.lkx.user;

import okhttp3.*;
import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;
import java.io.IOException;
import java.util.concurrent.TimeUnit;

public class ChatGPTClient {
    private static final String API_KEY = "sk-proj-ULBoqG88dqngMhEk8Su7T3BlbkFJdkOSvqPOkCiFUMtLJ7fp";
    private static final String API_URL = "https://www.cswx.net/v1/chat/completions";

    public static void main(String[] args) {
        try {
            String prompt = "今天是多少年多少月多少日";
            String response = getChatGPTResponse(prompt);
            System.out.println("ChatGPT response: " + response);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    public static String getChatGPTResponse(String prompt) throws IOException {
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .readTimeout(30, TimeUnit.SECONDS)
                .build();

        JSONObject requestBody = new JSONObject();
        try {
            requestBody.put("model", "gpt-4");
            requestBody.put("messages", new JSONArray().put(new JSONObject().put("role", "user").put("content", prompt)));
        } catch (JSONException e) {
            throw new RuntimeException(e);
        }


        Request request = new Request.Builder()
                .url(API_URL)
                .post(RequestBody.create(MediaType.parse("application/json"), requestBody.toString()))
                .addHeader("Authorization", "Bearer " + API_KEY)
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code " + response);
            }

            JSONObject jsonResponse = null;
            try {
                jsonResponse = new JSONObject(response.body().string());
                JSONArray choices = jsonResponse.getJSONArray("choices");
                return choices.getJSONObject(0).getJSONObject("message").getString("content");
            } catch (JSONException e) {
                throw new RuntimeException(e);
            }

        }
    }
}